package com.xp.ai.filesummary;

import com.google.common.collect.Lists;
import com.xp.ai.util.ModelUtils;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;

import java.io.File;
import java.io.FileInputStream;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;

/**
 * 文档提取摘要
 * @author xp
 */
public class FileSummary {
    public static String parseWord(String filePath) throws Exception {
        try (FileInputStream fis = new FileInputStream(filePath);
             XWPFDocument document = new XWPFDocument(fis);
             XWPFWordExtractor extractor = new XWPFWordExtractor(document)) {
            return extractor.getText();
        }
    }

    public static String parsePdf(String filePath) throws Exception {
        try (PDDocument document = PDDocument.load(new File(filePath))) {
            PDFTextStripper stripper = new PDFTextStripper();
            return stripper.getText(document);
        }
    }
    public static void main(String[] args) throws Exception {
//        Path path = Paths.get("E:\\weixin\\WeChat_Files\\WeChat Files\\wxid_mvp9im2rn3wi22\\FileStorage\\File\\2025-03\\4 绩效考核管理制度V1.4.pdf");
//        Path path = Paths.get("E:\\weixin\\WeChat_Files\\WeChat Files\\wxid_mvp9im2rn3wi22\\FileStorage\\File\\2025-03\\本周教科院直播课.docx");
//        smallFileParse();
        bigFileParse();
    }

    /**
     * 大文件解析并分块生成摘要
     * @throws Exception
     */
    private static void bigFileParse() throws Exception {
        Path path = Paths.get("E:\\weixin\\WeChat_Files\\WeChat Files\\wxid_mvp9im2rn3wi22\\FileStorage\\File\\2025-03\\校局协同育人平台建设方案0311(1).docx");
        //文档加载器
        Document document = FileSystemDocumentLoader.loadDocument(path, DocumentParseUtil.getParser(path.getFileName().toString()));
        //文档拆分器
        // 2. 分块处理
        DocumentSplitter splitter = DocumentSplitters.recursive(3000, 600);
        //文字块拆分
        List<TextSegment> segments = splitter.split(document);
        //将当前的文档内容进行向量化 embedding
        Thread.Builder.OfVirtual embeddingThread = Thread.ofVirtual().name("embedding");
        embeddingThread.start(() -> {});

        ChatLanguageModel deepSeeV3Model = ModelUtils.getDeepSeeV3Model();
        StringBuffer stringBuilder = new StringBuffer();
        // 4. 生成摘要
        String promptTemplate = "请用中文总结以下内容,保持200字以内：\n${text}\n摘要：";
        ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
        List<Future<String>> futures = new ArrayList<>();
        System.out.println("文本已经拆分为:"+ segments.size() +"块");
        for (TextSegment segment : segments) {
            String textToSummarize = segment.text();
            Future<String> submit = executorService.submit(() -> {
                String summary = ModelUtils.getDeepSeeV3Model().chat(promptTemplate.replace("${text}", textToSummarize));
                Thread.sleep(1500);
                System.out.println("摘要：" + summary);
                stringBuilder.append(summary).append("\n");
                return summary;
            });
            futures.add(submit);
        }
        futures.forEach(future -> {
            try {
                future.get();
            } catch (InterruptedException | ExecutionException e) {
                System.out.println(e);
            }
        });
        String summarizedText = deepSeeV3Model.chat(promptTemplate.replace("${text}", stringBuilder.toString()));
        System.out.println("最终摘要："+summarizedText);
        executorService.shutdown();

    }

    /**
     * 小文件解析
     * @throws Exception
     */
    private static void smallFileParse() throws Exception {
        Path path = Paths.get("E:\\weixin\\WeChat_Files\\WeChat Files\\wxid_mvp9im2rn3wi22\\FileStorage\\File\\2025-03\\校局协同育人平台建设方案0311(1).docx");
        //获取大模型
        ChatLanguageModel chatLanguageModel = ModelUtils.getHuoshanv3Model();
//        String text = parsePdf(path.toString());
        String text = parseWord(path.toString());
        //构建提示词模板
        String promptTemplate = "请用中文为以下文本生成摘要，要求保留关键信息，长度在200字以内：\n" + text;
        String summarizedText = chatLanguageModel.chat(promptTemplate);
        System.out.println(summarizedText);
    }
}
